import torch

flag = torch.cuda.is_available()
print(flag)

ngpu = 1
device = torch.device("cuda:0" if(torch.cuda.is_available() and ngpu > 0)else "cpu")

print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3, 3).cuda())

cuda_version = torch.version.cuda
print(cuda_version)

print(torch.backends.cudnn.version())